#ifndef GEO_MF_TOPIC_BACK_CC
#define GEO_MF_TOPIC_BACK_CC
#include<string>
#include<fstream>
#include<sstream>
#include<iostream>
#include<algorithm>
#include<iterator>
#include<vector>
#include<map>
#include<cstdlib>
#include<ctime>
#include<cmath>
#include "def.hpp"
#include "utils.hpp"
#include "topicmodel_geo_abstract.hpp"
#include "boost/unordered_map.hpp" // for hashtable

class GeoMFTopicBack:public GeoTopicModel{
protected:
  double* phi_0; // background language model for the whole corpus
  vector<int> doc_n;
  double** pi;
  double* phi_norm;
  double** theta_user; // background topic distribution for each user
  int* d_k;
  int** d_u_k;

  int**  n_k_v;
  int* n_v;
  int* n_k;
  int* d_u;

  int* doc_topic_z; // topic assignment for each document

  int N;
  int switch_label;
  double* theta_norm_buffer; // used for calculating theta norm
  double lowerbound;
  double learning_rate;
  double reg;
  double avg_term;
  double avg_change_term;
  int K; // parameter for L1 regularization by using Truncated Gradient method
  double g; // parameter for controling the sparsity
  double* p_k;
  ofstream file_output;
  int ITER;
public:
  GeoMFTopicBack();
  virtual void Inference();
  virtual void Prediction();
  virtual double Likelihood();
  virtual int ReadDOCS();
  virtual int GetDocuments(int); // this indicate how documents are read through into the model
  virtual int SaveModel();
private:
  void PhiNorm();
  void OptimizeThetaZero(int);
  void OptimizeThetaUser(int);
  void OptimizeThetaRegion(int);
  void OptimizePi(int);
  void OptimizePhiZero(int);
  int Single_M_Step();
  int Single_E_Step(int);
};

int GeoMFTopicBack::SaveModel(){
  ofstream outFile;
  string model_file = PuzaDEF::Instance()->output_file_name + ".pi";
  outFile.open(model_file.c_str());
  for (int i=0; i<PuzaDEF::Instance()->TERM_NUM; i++){
    outFile << i << "\t";
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      outFile << k << ":" << pi[k][i] << " ";
    }
    outFile << endl;
  }
  outFile.flush();
  outFile.close();

  /*  model_file = PuzaDEF::Instance()->output_file_name + ".phi0";
  outFile.open(model_file.c_str());
  for (int i=0; i<PuzaDEF::Instance()->TERM_NUM; i++){
    outFile << i << ":" << phi_0[i] << " " << endl;;
  }
  outFile.flush();
  outFile.close();
  */

  model_file = PuzaDEF::Instance()->output_file_name + ".theta_user";
  outFile.open(model_file.c_str());
  for(int i=0; i < PuzaDEF::Instance()->GEO_USER_NUM; i++){
    outFile << i << "\t";
    for (int k=0; k<PuzaDEF::Instance()->TOPIC_NUM; k++){
      outFile << k << ":" << theta_user[i][k] << " " ;
    }
    outFile << endl;
  }
  outFile.flush();
  outFile.close();

  model_file = PuzaDEF::Instance()->output_file_name + ".z";
  outFile.open(model_file.c_str());
  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int k = doc_topic_z[d];
    int region = doc_geo_stamp[d];
    int user = doc_user_stamp[d];
    outFile << d << "\t" << region << "\t" << user << "\t" << k << endl;
  }
  outFile.flush();
  outFile.close();
  return 0;
}

GeoMFTopicBack::GeoMFTopicBack(){
  allocate_memory<double>(theta_user, PuzaDEF::Instance()->GEO_USER_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<double>(pi, PuzaDEF::Instance()->TOPIC_NUM,PuzaDEF::Instance()->TERM_NUM);
  allocate_memory<int>(d_u_k, PuzaDEF::Instance()->GEO_USER_NUM,PuzaDEF::Instance()->TOPIC_NUM);
  allocate_memory<int>(n_k_v, PuzaDEF::Instance()->TOPIC_NUM,PuzaDEF::Instance()->TERM_NUM);

  phi_0= new double[PuzaDEF::Instance()->TERM_NUM];
  phi_norm = new double[PuzaDEF::Instance()->TOPIC_NUM];
  doc_topic_z = new int[PuzaDEF::Instance()->DOC_NUM];
  theta_norm_buffer = new double[PuzaDEF::Instance()->TOPIC_NUM];

  n_v = new int[PuzaDEF::Instance()->TERM_NUM];
  n_k = new int[PuzaDEF::Instance()->TOPIC_NUM];

  //  d_u_k.resize(PuzaDEF::Instance()->GEO_USER_NUM);
  d_k = new int[PuzaDEF::Instance()->TOPIC_NUM];
  p_k = new double[PuzaDEF::Instance()->TOPIC_NUM];
  d_u = new int[PuzaDEF::Instance()->GEO_USER_NUM];

  // some initialization
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k ++){
    d_k[k] = 0;
    n_k[k] = 0;
  }

  for(int u=0; u < PuzaDEF::Instance()->GEO_USER_NUM; u ++){
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      theta_user[u][k] = 0.0;
    }
    d_u[u] = 0;
  }

  for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
    phi_0[i] = 0.0;
  }

  learning_rate = PuzaDEF::Instance()->GEO_LEARNING_RATE;
  reg = PuzaDEF::Instance()->GEO_REG;
  K = PuzaDEF::Instance()->GEO_REG_L1_K;
  g = PuzaDEF::Instance()->GEO_REG_L1_G;
  avg_term = 0.0;
  avg_change_term = 0.0;
}

int GeoMFTopicBack::Single_E_Step(int d){
  // perform gibbs sampling for each document
  int new_topic = 0;
  int old_k = doc_topic_z[d];
  int user = doc_user_stamp[d];

  double u = 0.0;
  double theta_norm = 0.0;

  TermList current_list = TM_DOCS[d];
  int term_count = current_list.size();

  for(int k =0; k < PuzaDEF::Instance()->TOPIC_NUM;k++){
    theta_norm_buffer[k] = theta_user[user][k];
  }
  theta_norm = log_sum_array(theta_norm_buffer,PuzaDEF::Instance()->TOPIC_NUM); // calculate theta normalization part 

  for(int k =0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    p_k[k] = theta_user[user][k] - theta_norm;
    for (int j=0;j < term_count; j++){
      int term_id = current_list[j].term_id;
      int count = current_list[j].term_stat;
      double value = phi_0[term_id] + pi[k][term_id];
      p_k[k] = p_k[k] + count * (value - phi_norm[k]); // term level
    }
  }  

  for(int j=0; j < term_count; j++){
    int term_id = current_list[j].term_id;
    int count = current_list[j].term_stat;
    n_k_v[old_k][term_id] = n_k_v[old_k][term_id] - count;
  }
   
  // remove sufficient statistics
  d_k[old_k] --;
  d_u_k[user][old_k] --;
  n_k[old_k] -= doc_n[d];
  
  // convert back to probabilities and normalize
  p_k[0] = exp(p_k[0]);
  for(int k=1; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    p_k[k] = exp(p_k[k]) + p_k[k-1];
  }
  
  // sampling from uniform distribution
  u = gsl_rng_uniform(PuzaDEF::Instance()->gBaseRand) * p_k[PuzaDEF::Instance()->TOPIC_NUM-1];
  for(new_topic = 0; new_topic < PuzaDEF::Instance()->TOPIC_NUM; new_topic++){
    if (u < p_k[new_topic])
      break;
  }

  // add sufficient statistics
  n_k[new_topic] += doc_n[d];
  d_k[new_topic] ++;
  d_u_k[user][new_topic]++;
  for (int j=0;j < term_count; j++){
    int term_id = current_list[j].term_id;
    int count = current_list[j].term_stat;
    n_k_v[new_topic][term_id] += count;
  }

  if(new_topic != doc_topic_z[d]){
    switch_label ++;
    avg_change_term += doc_n[d];
  }

  avg_term += doc_n[d];

  doc_topic_z[d] = new_topic;  
  return 0;
}

void GeoMFTopicBack::OptimizeThetaUser(int shrink){
  for(int u=0; u < PuzaDEF::Instance()->GEO_USER_NUM;u++){
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      theta_norm_buffer[k] = theta_user[u][k];
    }
    double theta_norm = log_sum_array(theta_norm_buffer,PuzaDEF::Instance()->TOPIC_NUM);
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      double gradient = d_u_k[u][k]- d_u[u] * exp(theta_user[u][k] - theta_norm);
      theta_user[u][k] = theta_user[u][k] + learning_rate * gradient;
    }
  }
}

void GeoMFTopicBack::OptimizeThetaZero(int shrink){
  // update theta 0
  /*
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    double gradient = d_k[k] - d_k_gradient[k];
    theta_0[k] = theta_0[k] + learning_rate * gradient;
    d_k_gradient[k] = 0.0;
  }
  */
}

void GeoMFTopicBack::OptimizeThetaRegion(int shrink){
  // update theta geo
  /*  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM;k++){
      double gradient = 0.0;
      if(d_r_k[r].find(k) != d_r_k[r].end())
	gradient += d_r_k[r][k];
      gradient = gradient - d_r_gradient[r][k];
      theta_geo[r][k] = Trunc(theta_geo[r][k] + learning_rate * gradient, learning_rate * g);
      d_r_gradient[r][k] = 0.0;
    }
  }
  */
}

void GeoMFTopicBack::OptimizePi(int shrink){
  for(int k = 0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
      double gradient = n_k_v[k][v] - n_k[k] * exp(phi_0[v] + pi[k][v] - phi_norm[k]);
      pi[k][v] = pi[k][v] + learning_rate * gradient;
    }
  }
}

void GeoMFTopicBack::OptimizePhiZero(int shrink){
  
}

int GeoMFTopicBack::Single_M_Step(){
  int m_iter = 0;
  while(m_iter <= PuzaDEF::Instance()->GEO_M_ITER){
    OptimizePi(0);
    OptimizeThetaUser(0);  
    PhiNorm();
    m_iter ++;
  }
  return 0;
}

void GeoMFTopicBack::Inference(){
  PuzaLogger::Instance()->PutString("Start Inference for Geo Matrix Factorization Model [Only Topic Matrix]");
  if(PuzaDEF::Instance()->GEO_REG_TYPE == 0){
    PuzaLogger::Instance()->PutString("L2 Regularization.");
  }
  else{
    PuzaLogger::Instance()->PutString("L1 Regularization.");
  }
  ITER = 0;
  double old_lowerbound = 0.0;
  while(1){
    switch_label = 0;
    PuzaLogger::Instance()->PutString("Iteration "+to_string<int>(ITER,std::dec));
    for(unsigned int d=0; d < TM_DOCS.size(); d++){
      Single_E_Step(d);
    }
    PuzaLogger::Instance()->PutString("\tE-Step Finished."+to_string<int>(ITER,std::dec));
    Single_M_Step();
    PuzaLogger::Instance()->PutString("\tM-Step Finished."+to_string<int>(ITER,std::dec));
    Likelihood();
    PuzaLogger::Instance()->PutString("\tLabel Switch:" + to_string<int>(switch_label,std::dec)+" " + to_string<int>(ITER,std::dec));
    PuzaLogger::Instance()->PutString("\tAvg Change Terms:" + to_string<double>(avg_change_term/switch_label,std::dec)+" " + to_string<int>(ITER,std::dec));
    PuzaLogger::Instance()->PutString("\tAvg Terms:" + to_string<double>(avg_term/TM_DOCS.size(),std::dec)+" " + to_string<int>(ITER,std::dec));
    PuzaLogger::Instance()->PutString("\tLikelihood:" + to_string<double>(lowerbound,std::dec)+" " + to_string<int>(ITER,std::dec));

    avg_change_term = 0.0;
    avg_term = 0.0;

    old_lowerbound = lowerbound;
    if (ITER % 50 == 0){
      SaveModel();
    }

    if (ITER>=500){
      break;
    }

    ITER ++;
  }
  PuzaLogger::Instance()->PutString("Batch Learning Finished.");
}

void GeoMFTopicBack::Prediction(){
}
 
double GeoMFTopicBack::Likelihood(){
  double result = 0.0;

  for(unsigned int d = 0; d < TM_DOCS.size(); d++){
    int k = doc_topic_z[d];
    TermList current_list = TM_DOCS[d];
    int term_count = current_list.size();
    int user = doc_user_stamp[d];

    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      theta_norm_buffer[k] = theta_user[user][k];
    }
    double theta_norm = log_sum_array(theta_norm_buffer,PuzaDEF::Instance()->TOPIC_NUM);

    result = result + theta_user[user][k] - theta_norm;
    for (int j=0;j < term_count; j++){
      int term_id = current_list[j].term_id;
      int count = current_list[j].term_stat;
      double temp = phi_0[term_id] + pi[k][term_id];
      result = result + count * (temp - phi_norm[k]);
    }      
  }
  lowerbound = result;
  return lowerbound;
}

int GeoMFTopicBack::ReadDOCS(){
  int MAX_INT = 0;
  int c = 0;
  string temp;
  N = 0;
  if(PuzaDEF::Instance()->input_file_name.empty()){
    cout << "Please give input file name." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->InputFile.open(PuzaDEF::Instance()->input_file_name.c_str());
  if(!PuzaDEF::Instance()->InputFile){
    cout << "Can't open the file." << endl;
    exit(0);
  }
  PuzaDEF::Instance()->data_ids.clear();
  TM_DOCS.clear();
  doc_n.clear();
  string original_id;
  string temp_geo_stamp;
  string temp_user_stamp;
  vector<string> buffer;
  while(!PuzaDEF::Instance()->InputFile.eof()){
    getline(PuzaDEF::Instance()->InputFile,temp,'\n');
    std::istringstream iss(temp);
    // get the doc id
    getline(iss,temp,'\t');
    original_id = temp;
    // get geo                                                                                                                                 
    getline(iss,temp,'\t');
    temp_geo_stamp = temp;
    // get user
    getline(iss,temp,'\t');
    temp_user_stamp = temp;
    // get the doc body
    getline(iss,temp,'\n');
    std::istringstream tempiss(temp);
    vector<string> tokens;
    copy(istream_iterator<string>(tempiss), istream_iterator<string>(), back_inserter<vector<string> >(tokens));
    if(tokens.size()<1){
      continue;
    }

    int now_id = PuzaDEF::Instance()->data_ids.size();
    PuzaDEF::Instance()->data_ids[now_id] = original_id;
    boost::unordered_map<string,int>::iterator geo_iter = geo_reverse_mapping.find(temp_geo_stamp);
    int geo_id = -1;
    int user_id = -1;
    if(geo_iter != geo_reverse_mapping.end()){
      // id found
      geo_id = (*geo_iter).second;
      doc_geo_stamp.push_back(geo_id);
    }
    else{
      geo_id = geo_reverse_mapping.size();
      geo_reverse_mapping[temp_geo_stamp] = geo_id;
      geo_mapping[geo_id] = temp_geo_stamp;
      doc_geo_stamp.push_back(geo_id);
    }

    boost::unordered_map<string,int>::iterator user_iter = user_reverse_mapping.find(temp_user_stamp);
    if(user_iter != user_reverse_mapping.end()){
      // id found
      user_id = (*user_iter).second;
      doc_user_stamp.push_back(user_id);
    }
    else{
      user_id = user_reverse_mapping.size();
      user_reverse_mapping[temp_user_stamp] = user_id;
      user_mapping[user_id] = temp_user_stamp;
      doc_user_stamp.push_back(user_id);
    }
    // check user geo dist
    TermList assignList;
    int temp_n = 0;
    // update randomized parameters
    doc_topic_z[c] = (int)random_uniform_distribution(0, PuzaDEF::Instance()->TOPIC_NUM);
    int k = doc_topic_z[c];
    for (vector<string>::iterator iter = tokens.begin(); iter < tokens.end(); iter++){
      int word_id;
      int word_count;
      std::istringstream valueiss((*iter));
      getline(valueiss,temp,':');
      from_string<int>(word_id,temp,std::dec);
      getline(valueiss,temp);  // get count
      from_string<int>(word_count,temp,std::dec);
      TermItem newAssign;
      newAssign.term_id = word_id;
      newAssign.term_stat = word_count;
      assignList.push_back(newAssign);
      N = N + word_count;
      temp_n = temp_n + word_count;
      n_k_v[k][word_id] = n_k_v[k][word_id] + word_count;
      n_v[word_id] = n_v[word_id] + word_count;
    }
    if((int)assignList.size() > MAX_INT)
      MAX_INT = assignList.size();
    TM_DOCS.push_back(assignList);
    doc_n.push_back(temp_n);
    d_k[k] ++;
    d_u[user_id] ++;
    theta_user[user_id][k] ++;
    d_u_k[user_id][k] ++;
    n_k[k] += temp_n;
    c++;
  }
  PuzaDEF::Instance()->InputFile.close();
  return MAX_INT;
}

void GeoMFTopicBack::PhiNorm(){
  double* buffer = new double[PuzaDEF::Instance()->TERM_NUM];
  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
      buffer[i] = phi_0[i] + pi[k][i];
    }
    phi_norm[k] = log_sum_array(buffer,PuzaDEF::Instance()->TERM_NUM);
  }
  delete[] buffer;
}

int GeoMFTopicBack::GetDocuments(int choice){
  ReadDOCS();
  for(int v=0; v < PuzaDEF::Instance()->TERM_NUM; v++){
    phi_0[v] = log( (n_v[v] + 1.0) / (N + PuzaDEF::Instance()->TERM_NUM));
  }


  for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    for(int i=0; i < PuzaDEF::Instance()->TERM_NUM; i++){
      pi[k][i] = log( (n_k_v[k][i] + 1.0) / (n_k[k] + PuzaDEF::Instance()->TERM_NUM) ) - phi_0[i];
    }
  }

  PhiNorm();

  for(int u=0; u < PuzaDEF::Instance()->GEO_USER_NUM; u++){
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      theta_user[u][k] = log ((theta_user[u][k] + 1.0) / (d_u[u] + PuzaDEF::Instance()->TOPIC_NUM));
    }
  }

  // initilize theta parameters
  /*for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
    theta_0[k] = log( (d_k[k] + 1.0) / (TM_DOCS.size() + PuzaDEF::Instance()->TOPIC_NUM));
  }

  for(int r=0; r < PuzaDEF::Instance()->GEO_REGION_NUM; r++){
    for(int k=0; k < PuzaDEF::Instance()->TOPIC_NUM; k++){
      theta_geo[r][k] = log((theta_geo[r][k] + 1.0) / (d_r[r] + PuzaDEF::Instance()->TOPIC_NUM));
    }
  }
  */
  return 0;
}
#endif
